#include <iostream>
#include <array>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define MSG_LEN 8
#define MSG_CNT 100

#define AES_DIGEST_LEN ((16L - (MSG_LEN % 16)) + MSG_LEN)
#define INPUT_SIZE (MSG_LEN * MSG_CNT)
#define OUTPUT_SIZE (AES_DIGEST_LEN * MSG_CNT)
#define KEY_SIZE 32

enum class AES_TYPE {
  AES128 = 16,
  AES192 = 24,
  AES256 = 32
};

#ifndef WAFER_OPENSSL
extern "C"  {
    void CUDA(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt);
}

void AES_OP(MemRef<unsigned char, 1>* inputs, int msgLen, MemRef<unsigned char, 1>* outputs, int cptLen, MemRef<unsigned char, 1>* key, int keyLen, int msgCnt) {
    CUDA(inputs, msgLen, outputs, cptLen, key, keyLen, msgCnt);
}
#endif

uint8_t exampleInput[INPUT_SIZE] = {'0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9',
                                    '1', '0', '1', '1', '1', '2', '1', '3', '1', '4', '1', '5', '1', '6', '1', '7', '1', '8', '1', '9',
                                    '2', '0', '2', '1', '2', '2', '2', '3', '2', '4', '2', '5', '2', '6', '2', '7', '2', '8', '2', '9',
                                    '3', '0', '3', '1', '3', '2', '3', '3', '3', '4', '3', '5', '3', '6', '3', '7', '3', '8', '3', '9',
                                    '4', '0', '4', '1', '4', '2', '4', '3', '4', '4', '4', '5', '4', '6', '4', '7', '4', '8', '4', '9',
                                    '5', '0', '5', '1', '5', '2', '5', '3', '5', '4', '5', '5', '5', '6', '5', '7', '5', '8', '5', '9',
                                    '6', '0', '6', '1', '6', '2', '6', '3', '6', '4', '6', '5', '6', '6', '6', '7', '6', '8', '6', '9',
                                    '7', '0', '7', '1', '7', '2', '7', '3', '7', '4', '7', '5', '7', '6', '7', '7', '7', '8', '7', '9',
                                    '8', '0', '8', '1', '8', '2', '8', '3', '8', '4', '8', '5', '8', '6', '8', '7', '8', '8', '8', '9',
                                    '9', '0', '9', '1', '9', '2', '9', '3', '9', '4', '9', '5', '9', '6', '9', '7', '9', '8', '9', '9',
                                    '0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9',
                                    '1', '0', '1', '1', '1', '2', '1', '3', '1', '4', '1', '5', '1', '6', '1', '7', '1', '8', '1', '9',
                                    '2', '0', '2', '1', '2', '2', '2', '3', '2', '4', '2', '5', '2', '6', '2', '7', '2', '8', '2', '9',
                                    '3', '0', '3', '1', '3', '2', '3', '3', '3', '4', '3', '5', '3', '6', '3', '7', '3', '8', '3', '9',
                                    '4', '0', '4', '1', '4', '2', '4', '3', '4', '4', '4', '5', '4', '6', '4', '7', '4', '8', '4', '9',
                                    '5', '0', '5', '1', '5', '2', '5', '3', '5', '4', '5', '5', '5', '6', '5', '7', '5', '8', '5', '9',
                                    '6', '0', '6', '1', '6', '2', '6', '3', '6', '4', '6', '5', '6', '6', '6', '7', '6', '8', '6', '9',
                                    '7', '0', '7', '1', '7', '2', '7', '3', '7', '4', '7', '5', '7', '6', '7', '7', '7', '8', '7', '9',
                                    '8', '0', '8', '1', '8', '2', '8', '3', '8', '4', '8', '5', '8', '6', '8', '7', '8', '8', '8', '9',
                                    '9', '0', '9', '1', '9', '2', '9', '3', '9', '4', '9', '5', '9', '6', '9', '7', '9', '8', '9', '9',
                                    '0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9',
                                    '1', '0', '1', '1', '1', '2', '1', '3', '1', '4', '1', '5', '1', '6', '1', '7', '1', '8', '1', '9',
                                    '2', '0', '2', '1', '2', '2', '2', '3', '2', '4', '2', '5', '2', '6', '2', '7', '2', '8', '2', '9',
                                    '3', '0', '3', '1', '3', '2', '3', '3', '3', '4', '3', '5', '3', '6', '3', '7', '3', '8', '3', '9',
                                    '4', '0', '4', '1', '4', '2', '4', '3', '4', '4', '4', '5', '4', '6', '4', '7', '4', '8', '4', '9',
                                    '5', '0', '5', '1', '5', '2', '5', '3', '5', '4', '5', '5', '5', '6', '5', '7', '5', '8', '5', '9',
                                    '6', '0', '6', '1', '6', '2', '6', '3', '6', '4', '6', '5', '6', '6', '6', '7', '6', '8', '6', '9',
                                    '7', '0', '7', '1', '7', '2', '7', '3', '7', '4', '7', '5', '7', '6', '7', '7', '7', '8', '7', '9',
                                    '8', '0', '8', '1', '8', '2', '8', '3', '8', '4', '8', '5', '8', '6', '8', '7', '8', '8', '8', '9',
                                    '9', '0', '9', '1', '9', '2', '9', '3', '9', '4', '9', '5', '9', '6', '9', '7', '9', '8', '9', '9',
                                    '0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9',
                                    '1', '0', '1', '1', '1', '2', '1', '3', '1', '4', '1', '5', '1', '6', '1', '7', '1', '8', '1', '9',
                                    '2', '0', '2', '1', '2', '2', '2', '3', '2', '4', '2', '5', '2', '6', '2', '7', '2', '8', '2', '9',
                                    '3', '0', '3', '1', '3', '2', '3', '3', '3', '4', '3', '5', '3', '6', '3', '7', '3', '8', '3', '9',
                                    '4', '0', '4', '1', '4', '2', '4', '3', '4', '4', '4', '5', '4', '6', '4', '7', '4', '8', '4', '9',
                                    '5', '0', '5', '1', '5', '2', '5', '3', '5', '4', '5', '5', '5', '6', '5', '7', '5', '8', '5', '9',
                                    '6', '0', '6', '1', '6', '2', '6', '3', '6', '4', '6', '5', '6', '6', '6', '7', '6', '8', '6', '9',
                                    '7', '0', '7', '1', '7', '2', '7', '3', '7', '4', '7', '5', '7', '6', '7', '7', '7', '8', '7', '9',
                                    '8', '0', '8', '1', '8', '2', '8', '3', '8', '4', '8', '5', '8', '6', '8', '7', '8', '8', '8', '9',
                                    '9', '0', '9', '1', '9', '2', '9', '3', '9', '4', '9', '5', '9', '6', '9', '7', '9', '8', '9', '9'};
uint8_t exampleKey[KEY_SIZE] = {'A', '0', 'B', '1', 'C', '2', 'D', '3', 
                                'E', '4', 'F', '5', 'G', '6', 'H', '7', 
                                'I', '8', 'J', '9', 'K', '0', 'L', '1', 
                                'M', '2', 'N', '3', 'O', '4', 'P', '5',};

MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> key{exampleKey, KEY_SIZE};
uint8_t *exampleOutput = new uint8_t[OUTPUT_SIZE];
MemRef<unsigned char, 1> output{exampleOutput, OUTPUT_SIZE};

void printResult(uint8_t* output) {
    // Print the output.
    std::cout << "[ ";
    for (size_t i = 0; i < AES_DIGEST_LEN; i++) {
      printf("%02X", output[i]);
    }
    std::cout << "]" << std::endl;
}

void printInput(int idx, uint8_t* input) {
  // Print the input.
  std::cout << "Input " << idx << ": [";
  for (size_t i = 0; i < MSG_LEN; i++) {
    std::cout << input[i];
  }
  std::cout << "] ==> ";
}

void generateResultGPUAES() {
  int keyLen = static_cast<int>(AES_TYPE::AES128);

#ifdef WAFER_OPENSSL
  std::cout << "--------------------------------------------------------\nOpenSSL AES\n";

  const clock_t begin_time = clock();

  for (int i = 0; i < MSG_CNT; ++ i) {
    printInput(i, exampleInput + (i * MSG_LEN));

    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();

    const EVP_CIPHER *cipher;
    if (keyLen == 16)
      cipher = EVP_aes_128_ecb();
    else if (keyLen == 24)
      cipher = EVP_aes_192_ecb();
    else // keyLen == 32
      cipher = EVP_aes_256_ecb();

    EVP_EncryptInit_ex(ctx, cipher, NULL, exampleKey, NULL);
    int outlen;
    int ciphertext_len = 0;
    EVP_EncryptUpdate(ctx, exampleOutput + (i * AES_DIGEST_LEN), &outlen, exampleInput + (i * MSG_LEN), MSG_LEN);
    ciphertext_len += outlen;
    EVP_EncryptFinal_ex(ctx, exampleOutput + (i * AES_DIGEST_LEN) + ciphertext_len, &outlen);
    ciphertext_len += outlen;

    printResult(exampleOutput + (i * AES_DIGEST_LEN));
    EVP_CIPHER_CTX_free(ctx);
  }

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#else
  const clock_t begin_time = clock();
  #pragma wafer parallel
  AES_OP(&input, MSG_LEN, &output, AES_DIGEST_LEN, &key, keyLen, MSG_CNT);
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  
  std::cout << "\n--------------------------------------------------------\nMLIR AES\n";
  
  for (int i = 0; i < MSG_CNT; i++) {
    printInput(i, exampleInput + (i * MSG_LEN));
    printResult(exampleOutput + (i * AES_DIGEST_LEN));
  }

  std::cout << "Key " << ": [";
  for (size_t i = 0; i < keyLen; i++) {
    std::cout << key[i];
  }
  std::cout << "]" << std::endl;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

int main() {
  generateResultGPUAES();
  return 0;
}
